import torch
from torch import nn

class MLP(nn.Module):
    # 声明带有模型参数的层，这里声明了两个全连接层
    def __init__(self, n_feature):
        # 调用MLP父类Module的构造函数来进行必要的初始化。这样在构造实例时还可以指定其他函数
        # 参数，如“模型参数的访问、初始化和共享”一节将介绍的模型参数params
        super(MLP, self).__init__()
        self.hidden = nn.Linear(n_feature, 256) # 隐藏层
        self.hidden1 = nn.Linear(256, 128) # 隐藏层
        self.hidden2 = nn.Linear(128, 64) # 隐藏层
        self.hidden3 = nn.Linear(64, 32) # 隐藏层
        self.act = nn.ReLU()
        self.output = nn.Linear(32, 1)  # 输出层
         

    # 定义模型的前向计算，即如何根据输入x计算返回所需要的模型输出
    def forward(self, x):
        a = self.act(self.hidden(x))
        a=self.act(self.hidden1(a))
        a=self.act(self.hidden2(a))
        a=self.act(self.hidden3(a))
        return self.output(a)
# net=MLP(3)
# print(net.state_dict)
#print(net(torch.tensor([[1.0,2.0,3.0],[4.0,5.0,6]])))
#torch.save(net,'linear_model.pt')